package com.abc.util;

import cn.hutool.core.util.ObjectUtil;
import com.abc.config.MilvusProperties;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.DataType;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.FieldType;
import io.milvus.param.collection.HasCollectionParam;
import io.milvus.param.collection.LoadCollectionParam;
import io.milvus.param.index.CreateIndexParam;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;

/**
 * Milvus工具类
 * 提供向量数据库操作功能
 *
 * <p>过滤条件使用示例：</p>
 * <pre>
 * // 单条件过滤
 * Map&lt;String, Object&gt; filter1 = Map.of("category", "技术文档");
 *
 * // 多条件过滤（AND关系）
 * Map&lt;String, Object&gt; filter2 = Map.of(
 *     "category", "技术文档",
 *     "author", "张三",
 *     "status", "published"
 * );
 *
 * // 范围过滤
 * Map&lt;String, Object&gt; filter3 = Map.of(
 *     "createTime", Map.of("$gte", "2024-01-01", "$lte", "2024-12-31")
 * );
 *
 * // 使用过滤条件搜索
 * List&lt;Document&gt; results = milvusUtil.similaritySearch("Spring Boot", 5, filter1);
 * </pre>
 *
 * @author abc
 * @since 2024-01-01
 */
@Slf4j
@Component
public class MilvusUtil {

    @Autowired
    private MilvusServiceClient milvusClient;

    @Autowired
    private VectorStore vectorStore;

    @Autowired
    private MilvusProperties milvusProperties;

    /**
     * 检查集合是否存在
     */
    public boolean hasCollection(String databaseName, String collectionName) {
        try {
            HasCollectionParam param = HasCollectionParam.newBuilder()
                    .withCollectionName(collectionName)
                    .withDatabaseName(databaseName)
                    .build();
            return milvusClient.hasCollection(param).getData();
        } catch (Exception e) {
            log.error("检查集合是否存在失败: {}", collectionName, e);
            return false;
        }
    }

    /**
     * 创建集合
     */
    public boolean createCollection(String databaseName, String collectionName, int dimension) {
        try {
            if (hasCollection(databaseName, collectionName)) {
                log.info("集合已存在: {}", collectionName);
                return true;
            }

            FieldType idField = FieldType.newBuilder()
                    .withName("id")
                    .withDataType(DataType.Int64)
                    .withPrimaryKey(true)
                    .withAutoID(true)
                    .build();

            FieldType vectorField = FieldType.newBuilder()
                    .withName("embedding")
                    .withDataType(DataType.FloatVector)
                    .withDimension(dimension)
                    .build();

            FieldType contentField = FieldType.newBuilder()
                    .withName("content")
                    .withDataType(DataType.VarChar)
                    .withMaxLength(65535)
                    .build();

            CreateCollectionParam param = CreateCollectionParam.newBuilder()
                    .withCollectionName(collectionName)
                    .withDatabaseName(milvusProperties.getDatabaseName())
                    .withDescription("知识平台向量集合")
                    .withShardsNum(2)
                    .addFieldType(idField)
                    .addFieldType(vectorField)
                    .addFieldType(contentField)
                    .build();

            milvusClient.createCollection(param);
            log.info("集合创建成功: {}", collectionName);
            return true;
        } catch (Exception e) {
            log.error("创建集合失败: {}", collectionName, e);
            return false;
        }
    }

    /**
     * 创建索引
     */
    public boolean createIndex(String collectionName) {
        try {
            CreateIndexParam param = CreateIndexParam.newBuilder()
                    .withCollectionName(collectionName)
                    .withDatabaseName(milvusProperties.getDatabaseName())
                    .withFieldName("embedding")
                    .withIndexType(IndexType.valueOf(milvusProperties.getIndexType()))
                    .withMetricType(MetricType.valueOf(milvusProperties.getMetricType()))
                    .withExtraParam("{\"nlist\":1024}")
                    .build();

            milvusClient.createIndex(param);
            log.info("索引创建成功: {}", collectionName);
            return true;
        } catch (Exception e) {
            log.error("创建索引失败: {}", collectionName, e);
            return false;
        }
    }

    /**
     * 加载集合
     */
    public boolean loadCollection(String collectionName) {
        try {
            LoadCollectionParam param = LoadCollectionParam.newBuilder()
                    .withCollectionName(collectionName)
                    .withDatabaseName(milvusProperties.getDatabaseName())
                    .build();

            milvusClient.loadCollection(param);
            log.info("集合加载成功: {}", collectionName);
            return true;
        } catch (Exception e) {
            log.error("加载集合失败: {}", collectionName, e);
            return false;
        }
    }

    /**
     * 添加文档到向量存储
     */
    public void addDocuments(List<Document> documents) {
        try {
            if (ObjectUtil.isEmpty(documents)) {
                log.warn("文档列表为空，跳过添加操作");
                return;
            }

            vectorStore.add(documents);
            log.info("成功添加 {} 个文档到向量存储", documents.size());
        } catch (Exception e) {
            log.error("添加文档到向量存储失败", e);
            throw new RuntimeException("添加文档到向量存储失败", e);
        }
    }

    /**
     * 相似性搜索
     */
    public List<Document> similaritySearch(String query, int topK) {
        try {
            if (ObjectUtil.isEmpty(query)) {
                log.warn("查询内容为空");
                return List.of();
            }

            List<Document> results = vectorStore.similaritySearch(SearchRequest.builder().query(query).topK(topK).build());
            log.info("相似性搜索完成，查询: {}, 返回结果数: {}", query, results.size());
            return results;
        } catch (Exception e) {
            log.error("相似性搜索失败，查询: {}", query, e);
            throw new RuntimeException("相似性搜索失败", e);
        }
    }

    /**
     * 带过滤条件的相似性搜索
     *
     * @param query            查询文本
     * @param topK             返回结果数量
     * @param filterExpression 过滤条件，例如：
     *                         Map.of("category", "技术文档") - 按类别过滤
     *                         Map.of("author", "张三", "status", "published") - 多条件过滤
     *                         Map.of("createTime", Map.of("$gte", "2024-01-01")) - 时间范围过滤
     * @return 搜索结果文档列表
     */
    public List<Document> similaritySearch(String query, int topK, Map<String, Object> filterExpression) {
        try {
            if (ObjectUtil.isEmpty(query)) {
                log.warn("查询内容为空");
                return List.of();
            }

            SearchRequest.Builder searchRequestBuilder = SearchRequest.builder()
                    .query(query)
                    .topK(topK);

            // 如果有过滤条件，则添加到搜索请求中
            if (ObjectUtil.isNotEmpty(filterExpression)) {
                FilterExpressionBuilder filter = new FilterExpressionBuilder();
                filterExpression.forEach((key, value) -> {
                    searchRequestBuilder.filterExpression(filter.eq(key, value).build());
                });
                log.info("应用过滤条件: {}", filterExpression);
            }

            List<Document> results = vectorStore.similaritySearch(searchRequestBuilder.build());
            log.info("带过滤条件的相似性搜索完成，查询: {}, 过滤条件: {}, 返回结果数: {}",
                    query, filterExpression, results.size());
            return results;
        } catch (Exception e) {
            log.error("带过滤条件的相似性搜索失败，查询: {}, 过滤条件: {}", query, filterExpression, e);
            throw new RuntimeException("带过滤条件的相似性搜索失败", e);
        }
    }

    /**
     * 按单个条件过滤的相似性搜索
     *
     * @param query      查询文本
     * @param topK       返回结果数量
     * @param fieldName  过滤字段名
     * @param fieldValue 过滤字段值
     * @return 搜索结果文档列表
     */
    public List<Document> similaritySearchWithFilter(String query, int topK, String fieldName, Object fieldValue) {
        Map<String, Object> filterExpression = Map.of(fieldName, fieldValue);
        return similaritySearch(query, topK, filterExpression);
    }

    /**
     * 删除文档
     */
    public void deleteDocuments(List<String> ids) {
        try {
            if (ObjectUtil.isEmpty(ids)) {
                log.warn("文档ID列表为空，跳过删除操作");
                return;
            }

            vectorStore.delete(ids);
            log.info("成功删除 {} 个文档", ids.size());
        } catch (Exception e) {
            log.error("删除文档失败", e);
            throw new RuntimeException("删除文档失败", e);
        }
    }

    /**
     * 初始化向量存储
     */
    public void initializeVectorStore() {
        try {
            String collectionName = milvusProperties.getCollectionName();
            String databaseName = milvusProperties.getDatabaseName();
            int dimension = milvusProperties.getEmbeddingDimension();

            // 创建集合
            if (!createCollection(databaseName, collectionName, dimension)) {
                throw new RuntimeException("创建集合失败");
            }

            // 创建索引
            if (!createIndex(collectionName)) {
                throw new RuntimeException("创建索引失败");
            }

            // 加载集合
            if (!loadCollection(collectionName)) {
                throw new RuntimeException("加载集合失败");
            }

            log.info("向量存储初始化完成");
        } catch (Exception e) {
            log.error("向量存储初始化失败", e);
            throw new RuntimeException("向量存储初始化失败", e);
        }
    }
}

